import modelFit
from keras.callbacks import ModelCheckpoint
from flags import parse_args

#调用参数
FLAGS, unparsed = parse_args()

#训练数据
# train dataset
# load training dataset (6K)
filename = FLAGS.trainImages
train = modelFit.load_set(filename)
print('Dataset: %d' % len(train))
# descriptions
train_descriptions = modelFit.load_clean_descriptions(FLAGS.descriptions, train)
print('Descriptions: train=%d' % len(train_descriptions))
# photo features
train_features = modelFit.load_photo_features(FLAGS.features, train)
print('Photos: train=%d' % len(train_features))
# prepare tokenizer
tokenizer = modelFit.create_tokenizer(train_descriptions)
vocab_size = len(tokenizer.word_index) + 1
print('Vocabulary Size: %d' % vocab_size)
# determine the maximum sequence length
max_length = modelFit.max_length(train_descriptions)
print('Description Length: %d' % max_length)
# prepare sequences
X1train, X2train, ytrain = modelFit.create_sequences(
    tokenizer, max_length, train_descriptions, train_features,vocab_size)

# dev dataset

# load test set
filename = FLAGS.devImages
test = modelFit.load_set(filename)
print('Dataset: %d' % len(test))
# descriptions
test_descriptions = modelFit.load_clean_descriptions(FLAGS.descriptions, test)
print('Descriptions: test=%d' % len(test_descriptions))
# photo features
test_features = modelFit.load_photo_features(FLAGS.features, test)
print('Photos: test=%d' % len(test_features))
# prepare sequences
X1test, X2test, ytest = modelFit.create_sequences(
    tokenizer, max_length, test_descriptions, test_features,vocab_size)

# fit model

# define the model
model = modelFit.define_model(vocab_size, max_length,FLAGS.output_dir)
# define checkpoint callback
filepath = FLAGS.output_dir+'/model-ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'
checkpoint = ModelCheckpoint(
    filepath, monitor='val_loss', verbose=1, save_best_only=True, mode='min')
# fit model
model.fit([X1train, X2train], ytrain, epochs=20, verbose=2, callbacks=[
          checkpoint], validation_data=([X1test, X2test], ytest))